/**********************************************************************\

 Copyright 2012 Diggory Blake

 This file is part of klip-engine

 This program is free software: you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program.  If not, see <http://www.gnu.org/licenses/>.

\**********************************************************************/

#include "kimagefile.h"
#include "kerror.h"
#include "kdata.h"
#include "kstream.h"
#include <png.h>
#include <jpeglib.h>

template class kImageFileBase<kImageFilePNG>;
template class kImageFileBase<kImageFileJPG>;

std::map<keImageFormat, kImageFile*>& kImageFile::_fileTypes() {
	static std::map<keImageFormat, kImageFile*> fileTypes;
	return fileTypes;
}
bool kImageFile::load(kStream* source, kImageHeader* header, std::function<void (uint32_t, const uint8_t*)> lambda) {
	if (header->format == keImage_Unknown) {
		size_t pos = source->tell(keSeek_Read);
		for (auto i = _fileTypes().begin(); i != _fileTypes().end(); i++) {
			header->format = i->first;
			if (i->second->tryLoad(source, header, lambda)) {
				return true;
			} else {
				source->seek(keSeek_Read, pos);
			}
		}

		header->format = keImage_Unknown;
	} else {
		auto i = _fileTypes().find(header->format);
		if (i != _fileTypes().end()) {
			if (i->second->tryLoad(source, header, lambda)) {
				return true;
			}
		}
	}

	return false;
}
bool kImageFile::save(kStream* dest, const kImageHeader* header, std::function<void (uint32_t, uint8_t*)> lambda) {
	auto i = _fileTypes().find(header->format);
	if (i != _fileTypes().end()) {
		if (i->second->trySave(dest, header, lambda)) {
			return true;
		}
	}

	return false;
}
void png_error_func(png_struct* png_ptr, png_const_charp message) {
	throw kError(_T("Could not load/save png image!"));
}
void png_read_func(png_struct* png_ptr, png_byte* data, png_size_t length) {
	if (png_ptr == NULL)
		return;

	kStream* stream = (kStream*)png_get_io_ptr(png_ptr);

	if (!stream)
		png_error(png_ptr, "Internal read error");

	stream->read(kData(data, length));
}
void png_write_func(png_struct* png_ptr, png_byte* data, png_size_t length) {
	if (png_ptr == NULL)
		return;

	kStream* stream = (kStream*)png_get_io_ptr(png_ptr);

	if (!stream)
		png_error(png_ptr, "Internal read error");

	stream->write(kDataView(data, length));
}
bool kImageFilePNG::tryLoad(kStream* source, kImageHeader* header, std::function<void (uint32_t, const uint8_t*)> lambda) {
	png_struct* png_ptr = 0;
	png_info* info_ptr = 0;

	try {
		png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, &png_error_func, 0);
		info_ptr = png_create_info_struct(png_ptr);
	
		png_set_read_fn(png_ptr, source, &png_read_func);

		if (lambda) {
			png_read_png(png_ptr, info_ptr, PNG_TRANSFORM_STRIP_16|PNG_TRANSFORM_PACKING|PNG_TRANSFORM_EXPAND|PNG_TRANSFORM_GRAY_TO_RGB|PNG_TRANSFORM_BGR, 0);
		} else {
			png_read_info(png_ptr, info_ptr);
		}

		header->width = png_get_image_width(png_ptr, info_ptr);
		header->height = png_get_image_height(png_ptr, info_ptr);

		int fileDepth = (png_get_color_type(png_ptr, info_ptr) & PNG_COLOR_MASK_ALPHA) ? 4 : 3;

		if (header->depth != 3 && header->depth != 4)
			header->depth = fileDepth;
		
		uint8_t* row = 0;
		if (lambda) {
			if (header->depth != fileDepth) {
				row = (uint8_t*)malloc(header->width*header->depth);
			}
			png_byte** rows = png_get_rows(png_ptr, info_ptr);
			for (uint32_t y = 0; y < header->height; y++) {
				if (row) {
					for (uint32_t x = 0; x < header->width; x++) {
						row[x*header->depth] = rows[y][x*fileDepth];
						row[x*header->depth+1] = rows[y][x*fileDepth+1];
						row[x*header->depth+2] = rows[y][x*fileDepth+2];
						if (header->depth == 4)
							row[x*header->depth+3] = 0xFF;
					}
					lambda(y, row);
				} else {
					lambda(y, rows[y]);
				}
			}
			if (row)
				free(row);
		}
	} catch (...) {
		if (png_ptr)
			png_destroy_read_struct(&png_ptr, &info_ptr, NULL);

		return false;
	}

	return true;
}
bool kImageFilePNG::trySave(kStream* dest, const kImageHeader* header, std::function<void (uint32_t, uint8_t*)> lambda) {
	png_struct* png_ptr = 0;
	png_info* info_ptr = 0;

	try {
		png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, 0, &png_error_func, 0);
		info_ptr = png_create_info_struct(png_ptr);
	
		png_set_write_fn(png_ptr, dest, &png_write_func, 0);

		png_set_IHDR(png_ptr, info_ptr, header->width, header->height, 8,
			(header->depth == 3) ? PNG_COLOR_TYPE_RGB : PNG_COLOR_TYPE_RGBA,
			PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);

		uint8_t** rows = (uint8_t**)png_malloc(png_ptr, header->height*sizeof(uint8_t*));
		for (uint32_t y = 0; y < header->height; y++) {
			rows[y] = (uint8_t*)png_malloc(png_ptr, header->width*header->depth);
			lambda(y, rows[y]);
		}
		png_set_rows(png_ptr, info_ptr, rows);
		
		png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_PACKING|PNG_TRANSFORM_BGR, 0);
		
		for (uint32_t y = 0; y < header->height; y++) {
			png_free(png_ptr, rows[y]);
		}
		png_free(png_ptr, rows);

	} catch (...) {
		if (png_ptr)
			png_destroy_write_struct(&png_ptr, &info_ptr);

		return false;
	}

	return true;
}
struct jpg_src_mgr : jpeg_source_mgr {
	kStream* stream;
	kDataBuffer* buffer;

	inline jpg_src_mgr(kStream* source, size_t size) {
		stream = source;
		buffer = new kDataBuffer(size);
	}
	inline ~jpg_src_mgr() {
		delete buffer;
	}
};
void jpg_error_func(j_common_ptr cinfo) {
	throw kError(_T("Could not load/save jpg image!"));
}
boolean jpg_fill_input_buffer (j_decompress_ptr cinfo) {
    jpg_src_mgr* src = (jpg_src_mgr*)cinfo->src;
	src->bytes_in_buffer = src->stream->read(src->buffer->getData());
	src->next_input_byte = src->buffer->getPtr();

	return TRUE;
}
void jpg_init_source (j_decompress_ptr cinfo) {
	jpg_fill_input_buffer(cinfo);
}
void jpg_skip_input_data (j_decompress_ptr cinfo, long num_bytes)
{
	size_t nb = (size_t)num_bytes;
    jpg_src_mgr* src = (jpg_src_mgr*)cinfo->src;

    if (nb >= src->bytes_in_buffer) {
		nb -= src->bytes_in_buffer;

		if (nb)
			src->stream->skip(keSeek_Read, nb);

        jpg_fill_input_buffer(cinfo);
    } else {
		src->next_input_byte += nb;
		src->bytes_in_buffer -= nb;
	}
}
void jpg_term_source (j_decompress_ptr cinfo) {}
struct jpg_dest_mgr : jpeg_destination_mgr {
	kStream* stream;
	kDataBuffer* buffer;

	inline jpg_dest_mgr(kStream* dest, size_t size) {
		stream = dest;
		buffer = new kDataBuffer(size);
	}
	inline ~jpg_dest_mgr() {
		delete buffer;
	}
};
void jpg_init_dest (j_compress_ptr cinfo) {
	jpg_dest_mgr* dest = (jpg_dest_mgr*)cinfo->dest;
	dest->free_in_buffer = dest->buffer->getSize();
	dest->next_output_byte = dest->buffer->getPtr();
}
boolean jpg_empty_output_buffer (j_compress_ptr cinfo) {
	jpg_dest_mgr* dest = (jpg_dest_mgr*)cinfo->dest;
	dest->stream->write(dest->buffer->getDataView());
	dest->free_in_buffer = dest->buffer->getSize();
	dest->next_output_byte = dest->buffer->getPtr();

	return TRUE;
}
void jpg_term_destination (j_compress_ptr cinfo) {
	jpg_dest_mgr* dest = (jpg_dest_mgr*)cinfo->dest;
	dest->stream->write(kDataView(dest->buffer->getPtr(), dest->next_output_byte - dest->buffer->getPtr()));
}
bool kImageFileJPG::tryLoad(kStream* source, kImageHeader* header, std::function<void (uint32_t, const uint8_t*)> lambda) {
	jpeg_decompress_struct cinfo = {0};
	/* We use our private extension JPEG error handler.
	* Note that this struct must live as long as the main JPEG parameter
	* struct, to avoid dangling-pointer problems.
	*/
	jpeg_error_mgr jerr;
	/* More stuff */
	JSAMPARRAY buffer;		/* Output row buffer */
	int row_stride;		/* physical row width in output buffer */

	/* Step 1: allocate and initialize JPEG decompression object */

	/* We set up the normal JPEG error routines, then override error_exit. */
	cinfo.err = jpeg_std_error(&jerr);
	jerr.error_exit = &jpg_error_func;

	try {
		/* Now we can initialize the JPEG decompression object. */
		jpeg_create_decompress(&cinfo);

		/* Step 2: specify data source (eg, a file) */
		jpg_src_mgr src(source, 4096);

		src.init_source = &jpg_init_source;
		src.fill_input_buffer = &jpg_fill_input_buffer;
		src.skip_input_data = &jpg_skip_input_data;
		src.resync_to_restart = &jpeg_resync_to_restart; /* use default method */
		src.term_source = &jpg_term_source;
		src.bytes_in_buffer = 0;
		src.next_input_byte = 0;
		cinfo.src = &src;

		/* Step 3: read file parameters with jpeg_read_header() */
		jpeg_read_header(&cinfo, TRUE);

		/* Step 4: set parameters for decompression */

		header->width = cinfo.image_width;
		header->height = cinfo.image_height;
		if (!header->depth)
			header->depth = 3;

		if (lambda) {
			/* Step 5: Start decompressor */
			jpeg_start_decompress(&cinfo);

			/* We may need to do some setup of our own at this point before reading
			* the data.  After jpeg_start_decompress() we have the correct scaled
			* output image dimensions available, as well as the output colormap
			* if we asked for color quantization.
			* In this example, we need to make an output work buffer of the right size.
			*/ 
			/* JSAMPLEs per row in output buffer */
			row_stride = cinfo.output_width * cinfo.output_components;
			/* Make a one-row-high sample array that will go away when done with image */
			buffer = (*cinfo.mem->alloc_sarray)
			((j_common_ptr) &cinfo, JPOOL_IMAGE, row_stride, 1);
			uint8_t* outbuffer = (uint8_t*)malloc(header->width*header->depth);

			/* Step 6: while (scan lines remain to be read) */
			/*           jpeg_read_scanlines(...); */

			/* Here we use the library's state variable cinfo.output_scanline as the
			* loop counter, so that we don't have to keep track ourselves.
			*/
			while (cinfo.output_scanline < cinfo.output_height) {
				/* jpeg_read_scanlines expects an array of pointers to scanlines.
				* Here the array is only one element long, but you could ask for
				* more than one scanline at a time if that's more convenient.
				*/
				jpeg_read_scanlines(&cinfo, buffer, 1);

				for (uint32_t x = 0; x < header->width; x++) {
					outbuffer[x*header->depth] = buffer[0][x*cinfo.output_components+2];
					outbuffer[x*header->depth+1] = buffer[0][x*cinfo.output_components+1];
					outbuffer[x*header->depth+2] = buffer[0][x*cinfo.output_components];
					if (header->depth == 4)
						outbuffer[x*header->depth+3] = 0xFF;
				}

				/* Assume put_scanline_someplace wants a pointer and sample count. */
				lambda(cinfo.output_scanline-1, outbuffer);
			}

			free(outbuffer);

			/* Step 7: Finish decompression */

			jpeg_finish_decompress(&cinfo);
		}
		
		/* We can ignore the return value since suspension is not possible
		* with the stdio data source.
		*/

		/* Step 8: Release JPEG decompression object */

		/* This is an important step since it will release a good deal of memory. */
		jpeg_destroy_decompress(&cinfo);
	} catch (...) {
		jpeg_destroy_decompress(&cinfo);
		return false;
	}

	return true;
}
bool kImageFileJPG::trySave(kStream* dest, const kImageHeader* header, std::function<void (uint32_t, uint8_t*)> lambda) {
	jpeg_compress_struct cinfo;

	jpeg_error_mgr jerr;
	JSAMPROW row_pointer[1];	/* pointer to JSAMPLE row[s] */
	int row_stride;		/* physical row width in image buffer */

	/* Step 1: allocate and initialize JPEG compression object */

	/* We set up the normal JPEG error routines, then override error_exit. */
	cinfo.err = jpeg_std_error(&jerr);
	jerr.error_exit = &jpg_error_func;

	try {
		/* Now we can initialize the JPEG compression object. */
		jpeg_create_compress(&cinfo);

		/* Step 2: specify data destination (eg, a file) */
		/* Note: steps 2 and 3 can be done in either order. */

		jpg_dest_mgr dst(dest, 4096);
		dst.init_destination = &jpg_init_dest;
		dst.empty_output_buffer = &jpg_empty_output_buffer;
		dst.term_destination = &jpg_term_destination;
		cinfo.dest = &dst;

		/* Step 3: set parameters for compression */

		/* First we supply a description of the input image.
		* Four fields of the cinfo struct must be filled in:
		*/
		cinfo.image_width = header->width; 	/* image width and height, in pixels */
		cinfo.image_height = header->height;
		cinfo.input_components = 3;		/* # of color components per pixel */
		cinfo.in_color_space = JCS_RGB; 	/* colorspace of input image */
		/* Now use the library's routine to set default compression parameters.
		* (You must set at least cinfo.in_color_space before calling this,
		* since the defaults depend on the source color space.)
		*/
		jpeg_set_defaults(&cinfo);
		/* Now you can set any non-default parameters you wish to.
		* Here we just illustrate the use of quality (quantization table) scaling:
		*/
		jpeg_set_quality(&cinfo, 85, TRUE /* limit to baseline-JPEG values */);

		/* Step 4: Start compressor */

		/* TRUE ensures that we will write a complete interchange-JPEG file.
		* Pass TRUE unless you are very sure of what you're doing.
		*/
		jpeg_start_compress(&cinfo, TRUE);

		/* Step 5: while (scan lines remain to be written) */
		/*           jpeg_write_scanlines(...); */

		/* Here we use the library's state variable cinfo.next_scanline as the
		* loop counter, so that we don't have to keep track ourselves.
		* To keep things simple, we pass one scanline per call; you can pass
		* more if you wish, though.
		*/
		row_stride = header->width*header->depth;	/* JSAMPLEs per row in image_buffer */
		row_pointer[0] = (JSAMPROW)malloc(row_stride);

		while (cinfo.next_scanline < cinfo.image_height) {
			/* jpeg_write_scanlines expects an array of pointers to scanlines.
			* Here the array is only one element long, but you could pass
			* more than one scanline at a time if that's more convenient.
			*/
			lambda(cinfo.next_scanline, row_pointer[0]);

			if (header->depth == 4) {
				for (uint32_t x = 0; x < header->width; x++) {
					uint8_t tmp = row_pointer[0][x*4];
					row_pointer[0][x*3] = row_pointer[0][x*4+2];
					row_pointer[0][x*3+1] = row_pointer[0][x*4+1];
					row_pointer[0][x*3+2] = tmp;
				}
			}

			jpeg_write_scanlines(&cinfo, row_pointer, 1);
		}
		free(row_pointer[0]);

		/* Step 6: Finish compression */

		jpeg_finish_compress(&cinfo);

		/* Step 7: release JPEG compression object */

		/* This is an important step since it will release a good deal of memory. */
		jpeg_destroy_compress(&cinfo);

	} catch (...) {
		jpeg_destroy_compress(&cinfo);
		return false;
	}

	return true;
}
